import os
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QVBoxLayout, QTextEdit, QFileDialog, QMessageBox, QProgressBar
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QIcon
from PIL import Image
import pytesseract
from util import resource_path, set_window_geometry

class OcrApp(QWidget):
    def __init__(self):
        super().__init__()

        # 使用资源路径加载图标
        icon_path = resource_path('resource/app.png')
        self.setWindowIcon(QIcon(icon_path))  # 替换为你的图标文件路径
        # 设置窗口
        self.setWindowTitle("OCR 提取文字")
        # 使用封装的函数设置窗口的大小和位置
        set_window_geometry(self)

        # 创建布局
        layout = QVBoxLayout()

        # 创建显示提取文字的文本框
        self.text_edit = QTextEdit(self)
        self.text_edit.setReadOnly(True)
        layout.addWidget(self.text_edit)

        # 创建选择文件夹按钮
        self.select_folder_btn = QPushButton("选择图片文件夹", self)
        self.select_folder_btn.clicked.connect(self.select_folder)
        layout.addWidget(self.select_folder_btn)

        # 创建开始提取文字按钮
        self.extract_btn = QPushButton("开始提取文字", self)
        self.extract_btn.clicked.connect(self.extract_text)
        layout.addWidget(self.extract_btn)

        # 创建保存文件按钮
        self.save_btn = QPushButton("保存提取文本", self)
        self.save_btn.clicked.connect(self.save_text)
        layout.addWidget(self.save_btn)

        # 创建进度条
        self.progress_bar = QProgressBar(self)
        self.progress_bar.setRange(0, 100)  # 设置进度条的范围（0-100）
        self.progress_bar.setValue(0)  # 初始化进度条值为0
        layout.addWidget(self.progress_bar)

        # 设置布局
        self.setLayout(layout)

        # 初始化属性
        self.image_folder = ""
        self.extracted_text = ""

    def select_folder(self):
        """选择图片文件夹"""
        folder = QFileDialog.getExistingDirectory(self, "选择图片文件夹")
        if folder:
            self.image_folder = folder
            self.text_edit.setText(f"选择的文件夹: {folder}\n")

    def extract_text(self):
        """提取图片文件夹中的文字"""
        if not self.image_folder:
            self.text_edit.append("请先选择图片文件夹！")
            return

        # 获取图片文件夹中的所有图片文件
        image_files = [f for f in os.listdir(self.image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff'))]
        total_images = len(image_files)

        if total_images == 0:
            self.text_edit.append("该文件夹中没有图片文件！")
            return

        all_texts = []
        for index, image_name in enumerate(image_files):
            image_path = os.path.join(self.image_folder, image_name)

            try:
                # 打开图片
                img = Image.open(image_path)

                # 使用 pytesseract 识别文字
                text = pytesseract.image_to_string(img, lang='chi_sim')  # 使用中文识别
                all_texts.append(f"--- 文件: {image_name} ---\n{text}\n")

                self.text_edit.append(f"成功提取文字: {image_name}")
            except Exception as e:
                self.text_edit.append(f"处理文件 {image_name} 时出错: {e}")

            # 更新进度条
            progress = int((index + 1) / total_images * 100)
            self.progress_bar.setValue(progress)

        # 将提取的文字保存到属性中，以便后续保存
        self.extracted_text = "\n".join(all_texts)

        # 将提取的文字显示到文本框中
        if all_texts:
            self.text_edit.append("\n".join(all_texts))
        else:
            self.text_edit.append("未提取到任何文字。")

    def save_text(self):
        """将提取的文字保存到文件"""
        if not self.extracted_text:
            self.text_edit.append("没有提取到任何文字，无法保存！")
            return

        # 弹出文件保存对话框
        file_path, _ = QFileDialog.getSaveFileName(self, "保存文本", "", "Text Files (*.txt)")

        if file_path:
            try:
                # 保存提取的文本到文件
                with open(file_path, 'w', encoding='utf-8') as file:
                    file.write(self.extracted_text)
                self.text_edit.append(f"文本已保存至: {file_path}")
            except Exception as e:
                QMessageBox.warning(self, "保存失败", f"保存文件时出错: {e}")
        else:
            self.text_edit.append("保存操作已取消。")
